package cn.edu.recommender

import org.apache.spark.mllib.recommendation.{ALS, Rating}
import org.apache.spark.sql.SparkSession
import org.jblas.DoubleMatrix

import java.lang.Thread.sleep


case class ProductRating(userId: Int, productId: Int, score: Double, timestamp: Int)

case class MongoConfig(uri: String, db: String)

// 标准推荐对象，productId,score
case class Recommendation(productId: Int, score: Double)

// 用户推荐列表
case class UserRecs(userId: Int, recs: Seq[Recommendation])

// 商品相似度（商品推荐）
case class ProductRecs(productId: Int, recs: Seq[Recommendation])


object OfflineRecommender {


  // 定义常量
  val MONGODB_RATING_COLLECTION = "Rating"
  // 推荐表的名称
  val USER_RECS = "UserRecs"
  val PRODUCT_RECS = "ProductRecs"
  val USER_MAX_RECOMMENDATION = 20

  def main(args: Array[String]): Unit = {
    // 定义配置
    val config = Map(
      "spark.cores" -> "local[*]",
      "mongo.uri" -> "mongodb://localhost:27017/recommender",
      "mongo.db" -> "recommender"
    )

    val mongoConfig = MongoConfig(config("mongo.uri"), config("mongo.db"))

    // 创建一个 SparkSession
    val spark = SparkSession.builder().appName("StatisticsRecommender").master(config("spark.cores")).getOrCreate()

    // 在对 DataFrame 和 Dataset 进行操作许多操作都需要这个包进行支持
    import spark.implicits._

    //数据加载进来
    val ratingRDD = spark
      .read
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_RATING_COLLECTION)
      .format("com.mongodb.spark.sql")
      .load()
      .as[ProductRating]
      .map(rating => (rating.userId, rating.productId, rating.score))
      .rdd
      .cache()

    //用户与商品的数据集 RDD[Int]
    val userRDD = ratingRDD.map(_._1).distinct()
    val productRDD = ratingRDD.map(_._2).distinct()

    //创建训练数据集
    val trainData = ratingRDD.map(x => Rating(x._1, x._2, x._3))
    // rank 是模型中隐语义因子的个数, iterations 是迭代的次数, lambda 是 ALS 的正则化参
    val (rank, iterations, lambda) = (50, 5, 0.01)
    // 调用 ALS 算法训练隐语义模型
    val model = ALS.train(trainData, rank, iterations, lambda)

    // TODO：计算用户推荐矩阵
    val input = userRDD.cartesian(productRDD)
    // model 已训练好，把 userId,productId 传进去就可以得到预测评分列表 RDD[Rating] (userId,productId,rating)
    val output = model.predict(input)


    val userRecsDS = output
      .filter(_.rating > 0)
      .map(rating => (rating.user, (rating.product, rating.rating)))
      .groupByKey()
      .map {
        case (userId, recs) => UserRecs(userId, recs.toList.sortWith(_._2 > _._2).take(USER_MAX_RECOMMENDATION).map(x => Recommendation(x._1, x._2)))
      }.toDS()

    userRecsDS.write
      .option("uri", mongoConfig.uri)
      .option("collection", USER_RECS)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()

    // TODO：计算商品的兴趣相似度矩阵
    // 获取商品的特征矩阵，数据格式 RDD[(Int, Array[Double])]
    val productFeaturesRDD = model
      .productFeatures
      .map { case (productId, features) => (productId, new DoubleMatrix(features)) }

    // 计算笛卡尔积并过滤合并
    val productRecsDS = productFeaturesRDD.cartesian(productFeaturesRDD)
      .filter { case (a, b) => a._1 != b._1 }
      .map {
        case (a, b) =>
          val simScore = this.consinSim(a._2, b._2) // 求余弦相似度
          (a._1, (b._1, simScore))
      }
      .filter(_._2._2 > 0.6)
      .groupByKey()
      .map {
        case (productId, items) => ProductRecs(productId, items.toList.map(x => Recommendation(x._1, x._2)))
      }
      .toDS()


    productRecsDS
      .write
      .option("uri", mongoConfig.uri)
      .option("collection", PRODUCT_RECS)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()


    sleep(300000)
    // 关闭 Spark
    spark.stop()
  }

  //计算两个商品之间的余弦相似度
  def consinSim(product1: DoubleMatrix, product2:DoubleMatrix): Double = {
    product1.dot(product2) / ( product1.norm2() * product2.norm2() )
  }
}
